package com.ti.demo.controller;

import com.ti.demo.entity.ActorsFilms;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.converter.BeanOutputConverter;
import org.springframework.ai.converter.ListOutputConverter;
import org.springframework.ai.converter.MapOutputConverter;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.core.convert.support.DefaultConversionService;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

import java.util.List;
import java.util.Map;

/**
 * @description: com.ti.demo.controller
 * @Author Triagen
 * @Date 2025/7/18 10:53
 */
@RequestMapping("/structured")
@RestController
public class ChatController {

    private static final PromptTemplate BEAN_OUTPUT_PROMPT_TEMPLATE = new PromptTemplate("Generate the filmography of 5 movies for {actor}.");
    private static final String BEAN_OUTPUT_DIRECT_PROMPT = """
            Generate the filmography of 5 movies for {actor}.
            {format}
            """;
    private static final BeanOutputConverter<ActorsFilms> BEAN_OUTPUT_CONVERTER = new BeanOutputConverter<>(ActorsFilms.class);
    private static final BeanOutputConverter<List<ActorsFilms>> BEAN_OUTPUT_GENERIC_CONVERTER = new BeanOutputConverter<>(new ParameterizedTypeReference<>() {
    });



    private static final PromptTemplate MAP_OUTPUT_PROMPT_TEMPLATE = new PromptTemplate("Provide me a List of {subject}.");
    private static final String MAP_OUTPUT_DIRECT_PROMPT = """
            Provide me a List of {subject}.
            {format}
            """;
    private static final MapOutputConverter MAP_OUTPUT_CONVERTER = new MapOutputConverter();



    private static final PromptTemplate LIST_OUTPUT_PROMPT_TEMPLATE = new PromptTemplate("List five {subject}.");
    private static final String LIST_OUTPUT_DIRECT_PROMPT = """
            List five {subject}.
            {format}
            """;
    private static final ListOutputConverter LIST_OUTPUT_CONVERTER = new ListOutputConverter(new DefaultConversionService());

    @Qualifier("ollamaChatClient")
    @Autowired
    private ChatClient ollamaChatClient;


    @Autowired
    private ChatModel ollamaChatModel;


    @GetMapping(value = "/beanOutput")
    Mono<ActorsFilms> beanOutput(String actor) {
        return Mono.fromCallable(
                () -> this.ollamaChatClient.prompt()
                        .user(user -> user.text(BEAN_OUTPUT_PROMPT_TEMPLATE.getTemplate()).param("actor", actor))
                        .call()
                        .entity(ActorsFilms.class)
        ).subscribeOn(Schedulers.boundedElastic());
    }

    @GetMapping(value = "/beanOutputDirect")
    Mono<ActorsFilms> beanOutputDirect(String actor) {
        return Mono.fromCallable(() -> {
            String format = BEAN_OUTPUT_CONVERTER.getFormat();
            Generation generation = ollamaChatModel.call(
                    PromptTemplate.builder().template(BEAN_OUTPUT_DIRECT_PROMPT)
                            .variables(Map.of("format", format))
                            .build()
                            .create(Map.of("actor", actor))).getResult();
            return BEAN_OUTPUT_CONVERTER.convert(generation.getOutput().getText());
        }).subscribeOn(Schedulers.boundedElastic());

    }

    @GetMapping(value = "/beanOutputGeneric")
    Mono<List<ActorsFilms>> beanOutputGeneric(String actor) {
        return Mono.fromCallable(
                () -> this.ollamaChatClient.prompt()
                        .user(user -> user.text(BEAN_OUTPUT_PROMPT_TEMPLATE.getTemplate()).param("actor", actor))
                        .call()
                        .entity(new ParameterizedTypeReference<List<ActorsFilms>>() {
                        })
        ).subscribeOn(Schedulers.boundedElastic());
    }

    @GetMapping(value = "/beanOutputGenericDirect")
    Mono<List<ActorsFilms>> beanOutputGenericDirect(String actor) {
        return Mono.fromCallable(() -> {
            String format = BEAN_OUTPUT_GENERIC_CONVERTER.getFormat();
            Generation generation = ollamaChatModel.call(
                    PromptTemplate.builder().template(BEAN_OUTPUT_DIRECT_PROMPT)
                            .variables(Map.of("format", format))
                            .build()
                            .create(Map.of("actor", actor))).getResult();
            return BEAN_OUTPUT_GENERIC_CONVERTER.convert(generation.getOutput().getText());
        }).subscribeOn(Schedulers.boundedElastic());
    }


    @GetMapping(value = "/mapOutput")
    Mono<Map<String, Object>> mapOutput(String subject) {
        return Mono.fromCallable(
                () -> this.ollamaChatClient.prompt()
                        .user(user -> user.text(MAP_OUTPUT_PROMPT_TEMPLATE.getTemplate()).param("subject", subject))
                        .call()
                        .entity(new ParameterizedTypeReference<Map<String, Object>>() {
                        })
        ).subscribeOn(Schedulers.boundedElastic());
    }

    @GetMapping(value = "/mapOutputDirect")
    Mono<Map<String, Object>> mapOutputDirect(String subject) {
        return Mono.fromCallable(() -> {
            String format = MAP_OUTPUT_CONVERTER.getFormat();
            Generation generation = ollamaChatModel.call(
                    PromptTemplate.builder().template(MAP_OUTPUT_DIRECT_PROMPT)
                            .variables(Map.of("format", format))
                            .build()
                            .create(Map.of("subject", subject))).getResult();
            return MAP_OUTPUT_CONVERTER.convert(generation.getOutput().getText());
        }).subscribeOn(Schedulers.boundedElastic());
    }


    @GetMapping(value = "/listOutput")
    Mono<List<String>> listOutput(String subject) {
        return Mono.fromCallable(
                () -> this.ollamaChatClient.prompt()
                        .user(user -> user.text(LIST_OUTPUT_PROMPT_TEMPLATE.getTemplate()).param("subject", subject))
                        .call()
                        .entity(LIST_OUTPUT_CONVERTER)
        ).subscribeOn(Schedulers.boundedElastic());
    }

    @GetMapping(value = "/listOutputDirect")
    Mono<List<String>> listOutputDirect(String subject) {
        return Mono.fromCallable(() -> {
            String format = LIST_OUTPUT_CONVERTER.getFormat();
            Generation generation = ollamaChatModel.call(
                    PromptTemplate.builder().template(LIST_OUTPUT_DIRECT_PROMPT)
                            .variables(Map.of("format", format))
                            .build()
                            .create(Map.of("subject", subject))).getResult();
            return LIST_OUTPUT_CONVERTER.convert(generation.getOutput().getText());
        }).subscribeOn(Schedulers.boundedElastic());
    }

}
